from pathlib import Path
import math
import json
import numpy as np
from qtpy import QtWidgets
from qtpy import QtGui
from qtpy import QtCore
from qtpy.QtCore import Qt
import cv2


class Line(object):
    def __init__(self, x1, y1, x2, y2):
        self._x1 = x1
        self._y1 = y1
        self._x2 = x2
        self._y2 = y2
        self._pt1 = (x1, y1)
        self._pt2 = (x2, y2)

    @property
    def x1(self):
        return self._x1

    @property
    def x2(self):
        return self._x2

    @property
    def y1(self):
        return self._y1

    @property
    def y2(self):
        return self._y2

    @property
    def pt1(self):
        return self._pt1

    @property
    def pt2(self):
        return self._pt2

    def toArray(self):
        return np.array([self.x1, self.y1, self.x2, self.y2])


class CobbAngleWindow(QtWidgets.QWidget):
    keyPointNum = 17
    selectedLineNum = 4
    def __init__(self):
        super(CobbAngleWindow, self).__init__()
        self.setWindowTitle('cobb angle')
        self.setWindowFlags(Qt.WindowMaximizeButtonHint | Qt.MSWindowsFixedSizeDialogHint | Qt.WindowCloseButtonHint)
        self.imageLabel = QtWidgets.QLabel(self)
        self.showLabel = QtWidgets.QLabel(self)
        self.showLabel.setText('cobb angle')
        self.showLabel.setAlignment(Qt.AlignCenter)
        self.showLabel.setAutoFillBackground(True)
        palette = QtGui.QPalette()
        palette.setColor(QtGui.QPalette.Window, Qt.gray)
        self.showLabel.setPalette(palette)
        self.cobbButton = QtWidgets.QPushButton("compute cobb angle", self)
        self.layout = QtWidgets.QVBoxLayout(self)
        self.layout.addWidget(self.cobbButton)
        self.layout.addWidget(self.showLabel)
        self.layout.addWidget(self.imageLabel)
        self.setLayout(self.layout)
        self.scale = 0.3
        self.resize_image = None
        self.piximage = None
        self.image = None
        self.selectthreshold = 100
        self.height = 0
        self.width = 0
        self.channel = 0
        self.mid_lines = []
        self.selected_lines = []
        self.filename:Path =None

    @staticmethod
    def sort_point(shapes):
        point_dicts = {}
        for s in shapes:
            sk = int(s['label'])
            if sk not in point_dicts:
                point_dicts[sk] = [np.array(s['points'][0])]
            else:
                point_dicts[sk].append(np.array(s['points'][0]))

        assert len(point_dicts) == CobbAngleWindow.keyPointNum

        sort_point = []
        for i in range(CobbAngleWindow.keyPointNum):
            sps = np.array(point_dicts[i+1])
            sps = sps[sps[:, 1].argsort()]
            up_p = sps[0: 2, :]
            down_p = sps[2:4, :]
            up_p = up_p[up_p[:, 0].argsort()]
            down_p = down_p[down_p[:, 0].argsort()]
            sort_point.append(up_p[0])
            sort_point.append(up_p[1])
            sort_point.append(down_p[0])
            sort_point.append(down_p[1])
        return sort_point

    def getResult(self, image, shapes, filename):
        # 弹窗显示图像,并标中心线
        self.image = image
        self.filename = Path(filename)
        self.mid_lines = []
        self.selected_lines = []
        self.showLabel.setText('cobb angle')
        #sort point
        point_list = self.sort_point(shapes)

        for j in range(int(len(point_list) / 4)):
            middle_left = ((point_list[j * 4 + 2] + point_list[j * 4]) / 2).astype(int)
            middle_right = ((point_list[j * 4 + 3] + point_list[j * 4 + 1]) / 2).astype(int)
            self.mid_lines.append(Line(middle_left[0], middle_left[1], middle_right[0], middle_right[1]))
            cv2.circle(self.image, tuple(middle_left), 3, (255, 255, 0), 6)
            cv2.circle(self.image, tuple(middle_right), 3, (255, 255, 0), 6)
            cv2.line(self.image, (middle_left[0], middle_left[1]), (middle_right[0], middle_right[1]), (0, 255, 255), 8)

        self.showpiximage()
        self.cobbButton.clicked.connect(self.AngleCalculation)

    def getline(self, curPos):
        assert (len(self.mid_lines) == CobbAngleWindow.keyPointNum)

        def islinesFullOrRedundance(line: np.ndarray):
            if len(self.selected_lines) >= CobbAngleWindow.selectedLineNum:
                return True
            if len(self.selected_lines) == 0:
                return False
            logical = line == self.selected_lines
            logical = np.array([logical[i, :].all() for i in range(logical.shape[0])])
            return logical.any()

        for i in range(len(self.mid_lines)):
            m1 = np.array(self.mid_lines[i].pt1)
            m2 = np.array(self.mid_lines[i].pt2)
            m1_r = (m1 * self.scale).astype(int)
            m2_r = (m2 * self.scale).astype(int)
            result = (curPos[1] - m1_r[1]) * (m2_r[0] - m1_r[0]) - (m2_r[1] - m1_r[1]) * (curPos[0] - m1_r[0])
            full = islinesFullOrRedundance(self.mid_lines[i].toArray()) #判断当前线段是否在selected_lines中
            if abs(result) <= self.selectthreshold and not full:
                self.showLabel.setText(f'Select the {i}-th centerline')
                cv2.line(self.image, (m1[0], m1[1]), (m2[0], m2[1]), (255, 0, 0), 8)
                self.showpiximage()
                self.selected_lines.append(self.mid_lines[i].toArray())
                break
            else:
                continue
        #print("getline: ",len(self.selected_lines))

    def releaseline(self, curPos):
        for i in range(len(self.mid_lines)):
            m1 = np.array(self.mid_lines[i].pt1)
            m2 = np.array(self.mid_lines[i].pt2)
            m1_r = (m1 * self.scale).astype(int)
            m2_r = (m2 * self.scale).astype(int)
            result = (curPos[1] - m1_r[1]) * (m2_r[0] - m1_r[0]) - (m2_r[1] - m1_r[1]) * (curPos[0] - m1_r[0])

            if abs(result) <= self.selectthreshold:
                self.showLabel.setText(f'Cancel the {i}-th centerline')
                cv2.line(self.image, (m1[0], m1[1]), (m2[0], m2[1]), (0, 255, 255), 8)
                self.showpiximage()

                logical = self.mid_lines[i].toArray() == self.selected_lines
                logical = np.array([logical[i, :].all() for i in range(logical.shape[0])])
                del_index = np.argwhere(logical==True)[0][0]

                if len(self.selected_lines) > 0:
                    del self.selected_lines[del_index]
                break
            else:
                continue

        #print("releaseline: ", len(self.selected_lines))

    def showpiximage(self):
        resize_image = cv2.resize(self.image, (int(self.image.shape[1] * self.scale),int(self.image.shape[0] * self.scale)))
        height, width, channel = resize_image.shape
        Qimage = QtGui.QImage(resize_image.data, width, height, 3 * width, QtGui.QImage.Format_RGB888)
        self.imageLabel.setGeometry(0, 0, width, height)
        self.imageLabel.setAlignment(Qt.AlignCenter)
        self.piximage = QtGui.QPixmap.fromImage(Qimage)
        self.imageLabel.setPixmap(self.piximage)

        cobb_file = self.filename.with_name(self.filename.stem + '_cobb.json')
        if cobb_file.exists():
            with open(str(cobb_file), 'r') as f:
                cobb_dict = json.load(f)
            self.showLabel.setText(str(cobb_dict))

    def AngleCalculation(self):
        # calculate the middle points distance
        if len(self.selected_lines) != CobbAngleWindow.selectedLineNum:
            choice = QtWidgets.QMessageBox.question(self, 'centerlines', 'Please Select four centerlines! ', QtWidgets.QMessageBox.Yes)
            if choice == QtWidgets.QMessageBox.Yes:
                return

        vec_m_s = []
        for j in range(len(self.selected_lines)):
            #distance = self.mid_p_s[j * 2 + 1] - self.mid_p_s[j * 2]
            distance = np.array([self.selected_lines[j][2]-self.selected_lines[j][0], self.selected_lines[j][3]-self.selected_lines[j][1]])
            vec_m_s.append(distance)

        mod_v_s = np.power(np.sum(np.array(vec_m_s) * np.array(vec_m_s), axis=1), 0.5)
        # cobb angle
        cob_angles = []
        for i in range(0, len(mod_v_s)-1):
            mod_v = mod_v_s[i:i + 2]
            vec_m = vec_m_s[i:i + 2]
            dot_v = np.dot(np.array(vec_m), np.array(vec_m).T)
            mod_v = mod_v[..., np.newaxis]
            angle = np.arccos(np.round(dot_v / np.dot(mod_v, mod_v.T), 8))
            maxt = np.max(angle, axis=0)
            pt = np.max(maxt)
            cob_angles.append(round(pt / math.pi * 180, 2))

        cobb_dict={'pt': cob_angles[0], 'mt':cob_angles[1], 'tl':cob_angles[2]}

        save_file = self.filename.with_name(self.filename.stem+'_cobb.json')
        with open(str(save_file), 'w') as f:
            json.dump(cobb_dict, f, ensure_ascii=False, indent=2)
        self.showLabel.setText(str(cobb_dict))

    def mousePressEvent(self, evt):
        if evt.buttons() == QtCore.Qt.LeftButton:
            curPos = [evt.pos().x() - 12, evt.pos().y() - 58]
            self.getline(curPos)
        if evt.buttons() == QtCore.Qt.RightButton:
            curPos = [evt.pos().x() - 12, evt.pos().y() - 58]
            self.releaseline(curPos)


def test_sort_point():
    filename = 'E:/medical/vertebra/2D/data/boostnet_labeldata/data/training/sunhl-1th-05-Jan-2017-170 A AP.json'
    with open(filename, 'r') as f:
        data = json.load(f)

    shape_keys = [
        'label',
        'points',
        'group_id',
        'shape_type',
        'flags',
    ]
    shapes = [
        dict(
            label=s['label'],
            points=s['points'],
            shape_type=s.get('shape_type', 'polygon'),
            flags=s.get('flags', {}),
            group_id=s.get('group_id'),
            other_data={
                k: v for k, v in s.items() if k not in shape_keys
            }
        )
        for s in data['shapes']
    ]

    point_list = CobbAngleWindow.sort_point(shapes)


if __name__ == '__main__':
    print("")
    #test_sort_point()
    selected_lines = [np.array([12,23,34,45]),np.array([2,23,34,45]),np.array([12,3,34,45])]
    def islinesFullOrRedundance(line: np.ndarray):
        if len(selected_lines) >= CobbAngleWindow.selectedLineNum:
            return True
        logical = line == selected_lines
        logical = np.array([logical[i,:].all() for i in range(logical.shape[0])])
        return logical.any()


    line = np.array([2,23,34,45])

    print(islinesFullOrRedundance(line))


